import asyncio
import base64
import json
import os
import uuid

import httpx
from openai import AsyncOpenAI

from a2a.client import A2ACardResolver
from a2a.types import (
    AgentCard,
    DataPart,
    Message,
    MessageSendConfiguration,
    MessageSendParams,
    Part,
    Task,
    TaskState,
    TextPart,
)

from .remote_agent_connection import RemoteAgentConnections, TaskUpdateCallback


class HostAgent:
    """DeepSeek-powered host agent that maintains A2A protocol compatibility.

    This agent directly uses DeepSeek API for decision making while properly
    implementing A2A protocol for remote agent communication.
    """

    def __init__(
        self,
        remote_agent_addresses: list[str],
        http_client: httpx.AsyncClient,
        task_callback: TaskUpdateCallback | None = None,
    ):
        self.task_callback = task_callback
        self.httpx_client = http_client
        self.remote_agent_connections: dict[str, RemoteAgentConnections] = {}
        self.cards: dict[str, AgentCard] = {}
        self.agents: str = ''
        
        # 初始化DeepSeek客户端
        api_key = os.getenv('DEEPSEEK_API_KEY')
        if not api_key:
            print("Warning: DEEPSEEK_API_KEY not set. Please set it to use DeepSeek API.")
            # 使用一个临时的key来避免初始化错误
            api_key = "dummy_key"
            
        self.deepseek_client = AsyncOpenAI(
            api_key=api_key,
            base_url=os.getenv('DEEPSEEK_API_BASE', 'https://api.deepseek.com/v1')
        )
        self.model_name = os.getenv('DEEPSEEK_MODEL_NAME', 'deepseek-chat')
        
        # 会话状态管理
        self.session_contexts: dict[str, dict] = {}
        
        # 保存API key状态
        self._has_valid_api_key = bool(os.getenv('DEEPSEEK_API_KEY'))
        
        loop = asyncio.get_running_loop()
        loop.create_task(
            self.init_remote_agent_addresses(remote_agent_addresses)
        )

    async def init_remote_agent_addresses(
        self, remote_agent_addresses: list[str]
    ):
        async with asyncio.TaskGroup() as task_group:
            for address in remote_agent_addresses:
                task_group.create_task(self.retrieve_card(address))

    async def retrieve_card(self, address: str):
        card_resolver = A2ACardResolver(self.httpx_client, address)
        card = await card_resolver.get_agent_card()
        self.register_agent_card(card)

    def register_agent_card(self, card: AgentCard):
        remote_connection = RemoteAgentConnections(self.httpx_client, card)
        self.remote_agent_connections[card.name] = remote_connection
        self.cards[card.name] = card
        agent_info = []
        for ra in self.list_remote_agents():
            agent_info.append(json.dumps(ra))
        self.agents = '\n'.join(agent_info)

    def list_remote_agents(self):
        """List the available remote agents you can use to delegate the task."""
        if not self.remote_agent_connections:
            return []

        remote_agent_info = []
        for card in self.cards.values():
            remote_agent_info.append(
                {'name': card.name, 'description': card.description}
            )
        return remote_agent_info

    async def process_message(self, user_message: str, context: dict = None) -> str:
        """处理用户消息，使用DeepSeek进行决策，通过A2A协议与远程代理通信"""
        try:
            # 检查API key是否设置
            if not self._has_valid_api_key or not os.getenv('DEEPSEEK_API_KEY') or os.getenv('DEEPSEEK_API_KEY') == 'dummy_key':
                return ("DeepSeek API key not configured. Please set the DEEPSEEK_API_KEY environment variable. "
                       "You can also update it through the UI settings.")
            # 提取会话信息
            conversation_id = context.get('conversation_id') if context else None
            task_id = context.get('task_id') if context else None
            message_id = context.get('message_id') if context else None
            
            # 获取或创建会话上下文
            session_context = self._get_or_create_session_context(conversation_id)
            
            # 构建系统提示，包含历史对话
            system_prompt = self._build_system_prompt(session_context)
            
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message}
            ]
            
            # 添加历史对话到上下文
            if 'history' in session_context:
                # 插入历史对话在系统消息和当前用户消息之间
                messages = [messages[0]] + session_context['history'] + [messages[1]]

            # 调用DeepSeek API进行智能决策
            response = await self.deepseek_client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=0.7,
                max_tokens=1000
            )

            decision = response.choices[0].message.content
            
            # 更新会话历史
            self._update_session_history(session_context, user_message, decision)
            
            # 分析是否需要委托给远程代理
            delegation_result = await self._analyze_and_delegate(
                decision, user_message, conversation_id, task_id, message_id
            )
            
            if delegation_result:
                return f"{decision}\n\n--- Agent Response ---\n{delegation_result}"
            
            return decision

        except Exception as e:
            return f"Error processing message: {str(e)}"

    def update_api_key(self, api_key: str):
        """更新DeepSeek API key"""
        if api_key and api_key.strip():
            os.environ['DEEPSEEK_API_KEY'] = api_key.strip()
            self._has_valid_api_key = True
            # 重新初始化客户端
            self.deepseek_client = AsyncOpenAI(
                api_key=api_key.strip(),
                base_url=os.getenv('DEEPSEEK_API_BASE', 'https://api.deepseek.com/v1')
            )
            return True
        return False

    def _get_or_create_session_context(self, conversation_id: str) -> dict:
        """获取或创建会话上下文"""
        if not conversation_id:
            conversation_id = str(uuid.uuid4())
            
        if conversation_id not in self.session_contexts:
            self.session_contexts[conversation_id] = {
                'conversation_id': conversation_id,
                'history': [],
                'active_tasks': {},
                'last_agent': None
            }
        
        return self.session_contexts[conversation_id]

    def _build_system_prompt(self, session_context: dict) -> str:
        """构建包含代理信息和会话状态的系统提示"""
        base_prompt = f"""You are an expert AI orchestrator that coordinates tasks between specialized remote agents.

Available Remote Agents:
{self.agents}

Your responsibilities:
1. Analyze user requests to understand their intent and requirements
2. Determine if a request can be handled by available remote agents
3. If suitable agents exist, decide to delegate the task with clear instructions
4. Coordinate multi-step workflows when needed
5. Provide helpful responses when no suitable agent is available

Guidelines for delegation:
- When you decide to delegate, clearly state which agent you're choosing and why
- Include phrases like "I'll forward this to [Agent Name]" or "Let me ask the [Agent Name] to help"
- Be specific about what you're asking each agent to do
- For currency/exchange/conversion requests, use the Currency Agent

Session Info:
- Conversation ID: {session_context.get('conversation_id', 'new')}
- Last Agent Used: {session_context.get('last_agent', 'none')}
- Active Tasks: {len(session_context.get('active_tasks', {}))}
"""
        return base_prompt

    def _update_session_history(self, session_context: dict, user_message: str, assistant_response: str):
        """更新会话历史记录"""
        history = session_context.get('history', [])
        
        # 添加用户消息和助手响应
        history.append({"role": "user", "content": user_message})
        history.append({"role": "assistant", "content": assistant_response})
        
        # 保持历史记录在合理长度内（最近10轮对话）
        if len(history) > 20:  # 10轮对话 = 20条消息
            history = history[-20:]
        
        session_context['history'] = history

    async def _analyze_and_delegate(self, decision: str, original_message: str, 
                                   conversation_id: str, task_id: str, message_id: str) -> str:
        """分析决策并执行委托"""
        # 检查是否提到了委托
        if not self._should_delegate(decision):
            return None
            
        # 提取目标代理和任务信息
        agent_name, task_message = self._extract_delegation_info(decision, original_message)
        
        if not agent_name or agent_name not in self.remote_agent_connections:
            return f"Agent '{agent_name}' not available. Available agents: {list(self.remote_agent_connections.keys())}"
        
        # 更新会话上下文
        session_context = self._get_or_create_session_context(conversation_id)
        session_context['last_agent'] = agent_name
        
        # 执行A2A协议通信
        return await self._send_message_via_a2a(
            agent_name, task_message, conversation_id, task_id, message_id
        )

    def _should_delegate(self, decision: str) -> bool:
        """判断是否需要委托给远程代理"""
        delegation_indicators = [
            "forward this to", "ask the", "let me check with", "i'll use",
            "currency agent", "exchange rate", "convert", "delegate",
            "send this to", "route to", "contact the"
        ]
        decision_lower = decision.lower()
        return any(indicator in decision_lower for indicator in delegation_indicators)

    def _extract_delegation_info(self, decision: str, original_message: str) -> tuple[str, str]:
        """从决策中提取代理名称和任务消息"""
        decision_lower = decision.lower()
        
        # 识别目标代理
        agent_name = None
        for name in self.remote_agent_connections.keys():
            if name.lower() in decision_lower:
                agent_name = name
                break
        
        # 如果没有明确提到代理名，但内容涉及货币，使用Currency Agent
        if not agent_name:
            currency_keywords = ["currency", "exchange", "convert", "rate", "usd", "eur", "gbp", "jpy"]
            if any(keyword in decision_lower or keyword in original_message.lower() 
                   for keyword in currency_keywords):
                agent_name = "Currency Agent"
        
        # 构造发送给代理的消息（包含原始用户请求和决策上下文）
        if agent_name:
            task_message = f"""User Request: {original_message}

Context from orchestrator: {decision}

Please process this request and provide a helpful response."""
        else:
            task_message = original_message
            
        return agent_name, task_message

    async def _send_message_via_a2a(self, agent_name: str, message: str, 
                                   conversation_id: str, task_id: str, message_id: str) -> str:
        """通过A2A协议向远程代理发送消息"""
        try:
            client = self.remote_agent_connections[agent_name]
            
            # 使用现有ID或生成新ID，保持会话连续性
            actual_conversation_id = conversation_id or str(uuid.uuid4())
            actual_task_id = task_id or str(uuid.uuid4())
            actual_message_id = message_id or str(uuid.uuid4())
            
            # 构造符合A2A协议的请求
            request = MessageSendParams(
                id=str(uuid.uuid4()),  # Request ID
                message=Message(
                    role='user',
                    parts=[Part(root=TextPart(text=message))],
                    messageId=actual_message_id,
                    contextId=actual_conversation_id,
                    taskId=actual_task_id,
                ),
                configuration=MessageSendConfiguration(
                    acceptedOutputModes=['text', 'text/plain'],
                ),
            )
            
            # 发送消息并处理响应
            response = await client.send_message(request, self.task_callback)
            
            # 处理不同类型的响应
            if isinstance(response, Message):
                # 直接消息响应
                return self._extract_text_from_message(response)
            elif isinstance(response, Task):
                # 任务响应
                return self._extract_text_from_task(response)
            else:
                return f"Received response from {agent_name}: {str(response)}"
                
        except Exception as e:
            return f"Error communicating with {agent_name}: {str(e)}"

    def _extract_text_from_message(self, message: Message) -> str:
        """从Message对象中提取文本内容"""
        if not message.parts:
            return "No content in response"
            
        text_parts = []
        for part in message.parts:
            if part.root.kind == 'text':
                text_parts.append(part.root.text)
            elif part.root.kind == 'data':
                text_parts.append(str(part.root.data))
            
        return '\n'.join(text_parts) if text_parts else "No text content in response"

    def _extract_text_from_task(self, task: Task) -> str:
        """从Task对象中提取文本内容"""
        text_parts = []
        
        # 从任务状态消息中提取
        if task.status and task.status.message:
            text_parts.append(self._extract_text_from_message(task.status.message))
        
        # 从任务工件中提取
        if task.artifacts:
            for artifact in task.artifacts:
                for part in artifact.parts:
                    if part.root.kind == 'text':
                        text_parts.append(part.root.text)
                        
        return '\n'.join(text_parts) if text_parts else f"Task {task.id} completed"

    # ADK兼容性接口
    def create_agent(self):
        """返回自身，保持ADK兼容性"""
        return self

    @property 
    def name(self):
        return "deepseek_host_agent"

    async def run_async(self, user_id: str, session_id: str, new_message, **kwargs):
        """ADK兼容性方法，处理来自ADK Runner的调用"""
        if hasattr(new_message, 'parts') and new_message.parts:
            message_text = new_message.parts[0].text
        else:
            message_text = str(new_message)
        
        context = {
            'conversation_id': session_id,
            'user_id': user_id
        }
        
        result = await self.process_message(message_text, context)
        
        # 返回ADK兼容的事件对象
        class SimpleEvent:
            def __init__(self, content):
                self.content = SimpleContent(content)
                self.is_final = True
            
            def is_final_response(self):
                return self.is_final
        
        class SimpleContent:
            def __init__(self, text):
                self.parts = [SimplePart(text)]
                self.role = 'agent'
        
        class SimplePart:
            def __init__(self, text):
                self.text = text
        
        yield SimpleEvent(result)


async def convert_parts(parts: list[Part], tool_context=None):
    rval = []
    for p in parts:
        rval.append(await convert_part(p, tool_context))
    return rval


async def convert_part(part: Part, tool_context=None):
    if part.root.kind == 'text':
        return part.root.text
    if part.root.kind == 'data':
        return part.root.data
    if part.root.kind == 'file':
        return f"File: {part.root.file.name}"
    return f'Unknown type: {part.kind}'
